import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeRegressor
import matplotlib.pyplot as plt
from sklearn.preprocessing import LabelEncoder

# 西瓜数据集
dataSet = [['青绿', '蜷缩', '浊响', '清晰', '凹陷', '硬滑', '好瓜'], 
          ['乌黑', '蜷缩', '沉闷', '清晰', '凹陷', '硬滑', '好瓜'], 
          ['乌黑', '蜷缩', '浊响', '清晰', '凹陷', '硬滑', '好瓜'], 
          ['青绿', '蜷缩', '沉闷', '清晰', '凹陷', '硬滑', '好瓜'], 
          ['浅白', '蜷缩', '浊响', '清晰', '凹陷', '硬滑', '好瓜'], 
          ['青绿', '稍蜷', '浊响', '清晰', '稍凹', '软粘', '好瓜'], 
          ['乌黑', '稍蜷', '浊响', '稍糊', '稍凹', '软粘', '好瓜'], 
          ['乌黑', '稍蜷', '浊响', '清晰', '稍凹', '硬滑', '好瓜'], 
          ['乌黑', '稍蜷', '沉闷', '稍糊', '稍凹', '硬滑', '坏瓜'], 
          ['青绿', '硬挺', '清脆', '清晰', '平坦', '软粘', '坏瓜'], 
          ['浅白', '硬挺', '清脆', '模糊', '平坦', '硬滑', '坏瓜'], 
          ['浅白', '蜷缩', '浊响', '模糊', '平坦', '软粘', '坏瓜'], 
          ['青绿', '稍蜷', '浊响', '稍糊', '凹陷', '硬滑', '坏瓜'], 
          ['浅白', '稍蜷', '沉闷', '稍糊', '凹陷', '硬滑', '坏瓜'], 
          ['乌黑', '稍蜷', '浊响', '清晰', '稍凹', '软粘', '坏瓜'], 
          ['浅白', '蜷缩', '浊响', '模糊', '平坦', '硬滑', '坏瓜'], 
          ['青绿', '蜷缩', '沉闷', '稍糊', '稍凹', '硬滑', '坏瓜']]

# 特征名称
feature_names = ['色泽', '根蒂', '敲声', '纹理', '脐部', '触感', '好坏']

def prepareDataSet():
    """准备数据集，将文本特征转换为数值特征"""
    # 创建DataFrame
    df = pd.DataFrame(dataSet, columns=feature_names)
    
    # 分离特征和标签
    X = df.iloc[:, :-1]  # 所有行，除最后一列外的所有列作为特征
    y = df.iloc[:, -1]   # 所有行，最后一列作为标签
    
    # 对标签进行编码：好瓜->1，坏瓜->0
    label_encoder = LabelEncoder()
    y_encoded = label_encoder.fit_transform(y)
    
    # 对特征进行编码
    X_encoded = X.copy()
    for column in X.columns:
        X_encoded[column] = label_encoder.fit_transform(X[column])
    
    return X_encoded, y_encoded

def getTrainSetAndTestSet(X, y):
    """划分训练集和测试集"""
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=42)
    print("训练集特征维度:", X_train.shape)
    print("训练集标签维度:", y_train.shape)
    print("测试集特征维度:", X_test.shape)
    print("测试集标签维度:", y_test.shape)
    return X_train, X_test, y_train, y_test

def TrainDecisionTreeRegressor(X_train, y_train):
    """训练决策树回归模型"""
    # 创建决策树回归模型
    tree_reg = DecisionTreeRegressor(criterion='squared_error', max_depth=3, random_state=1)
    # 训练模型
    tree_reg.fit(X_train, y_train)
    print("决策树特征重要性:", tree_reg.feature_importances_)
    return tree_reg

def EvaluationModel(tree_reg, X_test, y_test):
    """评估模型性能"""
    y_pred = tree_reg.predict(X_test)
    mse = np.mean((y_pred - y_test) ** 2)
    print("均方误差MSE:", mse)
    rmse = np.sqrt(mse)
    print("均方根误差RMSE:", rmse)
    
    # 计算准确率（将预测值四舍五入后与实际值比较）
    y_pred_rounded = np.round(y_pred)
    accuracy = np.mean(y_pred_rounded == y_test)
    print("准确率:", accuracy)
    
    return y_pred

def Visualization(y_test, y_pred):
    """可视化结果"""
    fig, ax = plt.subplots()
    ax.scatter(y_test, y_pred)
    ax.plot([min(y_test), max(y_test)], [min(y_test), max(y_test)], 'k--', lw=5)
    ax.set_xlabel("Measured")
    ax.set_ylabel("Predicted")
    
    plt.show()
    plt.close(fig)

if __name__ == "__main__":
    # 准备数据集
    X, y = prepareDataSet()
    
    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = getTrainSetAndTestSet(X, y)
    
    # 训练决策树回归模型
    tree_reg_model = TrainDecisionTreeRegressor(X_train, y_train)
    
    # 评估模型性能
    y_pred = EvaluationModel(tree_reg_model, X_test, y_test)
    
    # 可视化结果
    Visualization(y_test, y_pred)